package com.liaoyifan.core.datasource;

import com.liaoyifan.core.converter.HikariConverter;
import com.liaoyifan.core.properties.DataSourceProperties;
import com.liaoyifan.core.properties.HikariProperties;
import com.liaoyifan.core.valid.Assert;
import com.liaoyifan.core.valid.Lang;
import com.zaxxer.hikari.HikariConfig;
import com.zaxxer.hikari.HikariDataSource;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import javax.sql.DataSource;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

public class DynamicDataSource extends AbstractRoutingDataSource {

    private final String primary;
    private final Map<String, HikariProperties> hikariProperties = new ConcurrentHashMap<>();
    private final Map<Object, Object> dynamicDataSources = new ConcurrentHashMap<>();

    public DynamicDataSource(DataSourceProperties dataSourceProperties) {
        this.hikariProperties.putAll(dataSourceProperties.getHikari());
        Assert.notEmpty(this.hikariProperties, "DataSource properties cannot be empty");
        this.primary =
                this.hikariProperties.size() == 1
                        ? hikariProperties.entrySet().iterator().next().getKey()
                        : dataSourceProperties.getPrimary();
        Assert.notBlank(this.primary, "DataSource primary cannot be blank");
        DataSource defaultDataSource = this.createDataSource(this.primary);
        this.dynamicDataSources.put(this.primary, defaultDataSource);
        super.setDefaultTargetDataSource(defaultDataSource);
        this.refreshDataSources();
    }

    @Override
    protected Object determineCurrentLookupKey() {
        String current = DataSourceHolder.get();
        String key = Lang.isBlank(current) ? this.primary : current;
        boolean refresh = !this.dynamicDataSources.containsKey(key);
        this.dynamicDataSources.computeIfAbsent(key, i -> createDataSource((String) i));
        if (refresh) {
            this.refreshDataSources();
        }
        return key;
    }

    private synchronized void refreshDataSources() {
        super.setTargetDataSources(this.dynamicDataSources);
        super.afterPropertiesSet();
    }

    private DataSource createDataSource(String key) {
        HikariProperties properties = this.hikariProperties.get(key);
        Assert.notNull(properties, "DataSource properties not found for key: " + key);
        HikariConfig config = HikariConverter.INSTANCE.convert(properties);
        Assert.notBlank(config.getJdbcUrl(), "DataSource jdbc url is required");
        config.setPoolName(key + "Pool");
        return new HikariDataSource(config);
    }
}
